/* DHCP Client */

/* Copyright (C) 2011 David Zanetti
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; version 2 of the License.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with this program; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include <stdio.h>
#include <stdlib.h>
#include <avr/io.h>
#include <util/atomic.h>
#include "global.h"
#include <util/delay.h>
#include <string.h>
#include "encx24j600.h"
#include "serial.h"
#include <avr/pgmspace.h>
#include "dhcpc.h"
#include "clock.h"
#include "udp.h"
#include "ip.h"

/* states as per RFC2131, we don't use INIT-REBOOT or REBOOTING */
#define DHCPC_INIT 0
#define DHCPC_SELECTING 1
#define DHCPC_REQUESTING 2
#define DHCPC_BOUND 3
#define DHCPC_RENEWING 4
#define DHCPC_REBINDING 5

/* there are other states but we don't care too much about them */
uint8_t state;

/* DHCP mesaage types for option 53 */

#define DHCPDISCOVER 1
#define DHCPOFFER 2
#define DHCPREQUEST 3
#define DHCPDECLINE 4
#define DHCPACK 5
#define DHCPNAK 6
#define DHCPRELEASE 7
#define DHCPINFORM 8

/* BOOTP opcodes */
#define BOOTREQUEST 1
#define BOOTREPLY 2

/* DHCP packet for DHCPDISCOVER */
typedef struct {
		uint8_t op;
		uint8_t htype;
		uint8_t hlen;
		uint8_t hops;
		uint8_t xid[4];
		uint8_t secs[2];
		uint8_t flags[2];
		uint8_t ciaddr[4];
		uint8_t yiaddr[4];
		uint8_t siaddr[4];
		uint8_t giaddr[4];
		uint8_t chaddr[16];
		uint8_t sname[64]; /* this could contain options too, thanks guys */
		uint8_t file[128];
		uint8_t magic[4]; /* DHCP packets *always* have this, bootp not so much */
		uint8_t option_code; /* DHCP packets also must always begin with option 53 */
		uint8_t option_len;
		uint8_t option_type;
		uint8_t end; /* the end options byte */
} dhcp_tx_header_t;

/* since we don't know the order of the options, RX side needs to process them seperately */
typedef struct {
		uint8_t op;
		uint8_t htype;
		uint8_t hlen;
		uint8_t hops;
		uint8_t xid[4];
		uint8_t secs[2];
		uint8_t flags[2];
		uint8_t ciaddr[4];
		uint8_t yiaddr[4];
		uint8_t siaddr[4];
		uint8_t giaddr[4];
		uint8_t chaddr[16];
		uint8_t sname[64]; /* this could contain options too, thanks guys */
		uint8_t file[128];
		uint8_t magic[4]; /* DHCP packets *always* have this, bootp not so much */
} dhcp_rx_header_t;

uint32_t last_xid; /* the last xid value we sent out. tick updates this for DHCPDISCOVER, and we discard replies without the "current" one */

uint32_t timer = 0; /* a timer for events. see FLAG_IP_DHCP_TIMER */
uint32_t start = 0; /* when we started either the discovery or renewal */
uint16_t secs = 0; /* DHCPREQUEST needs the same secs as our discover had */
uint32_t renewal = 0; /* how long before we must renew */
uint32_t rebind = 0; /* how long before we must rebind */
uint8_t discoverwait = 2; /* incrementing timeout for DHCPDISCOVER */

uint8_t yiaddr[4]; /* our address before we have committed it */

void dhcpc_init(void) {
	/* start in the init state, tick will take care of sendins something
    useful in this state */
	state = DHCPC_INIT;
	/* set the timer to 0 but let it run so we send DHCPDISCOVER */
	timer = 0;
	set_flag(flag_ip,FLAG_IP_DHCP_TIMER);
	return;
}

/* process some packet coming to us on port udp/68 */
uint8_t dhcpc_process(uint8_t slot) {
	dhcp_rx_header_t pkt;
	uint8_t dhcp_type = 0;
	uint8_t dhcp_si[4] = {0,0,0,0};
	uint8_t dhcp_mask[4] = {0,0,0,0};
	uint8_t dhcp_gw[4] = {0,0,0,0};
	uint32_t dhcp_renewal = 0;
	uint32_t dhcp_rebind = 0;
	uint8_t option_code, option_len;
	uint8_t nomoreoptions = 0;

	/* read the DHCP header at least, should contain more options, but
    we'll deal with those shortly */
	if (encx24j600_rx(&pkt,sizeof(dhcp_rx_header_t)) != sizeof(dhcp_rx_header_t)) {
		serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [bootp header]\r\n"));
		return E_NOTHING;
	}

	if (pkt.magic[0] != 0x63 ||
			pkt.magic[1] != 0x82 ||
			pkt.magic[2] != 0x53 ||
			pkt.magic[3] != 0x63) {
		/* magic cookie FAILIED */
		serial0_tx_PGM(PSTR("dhcpc: bad magic cookie, discarding\r\n"));
		return E_NOTHING;
	}
#ifdef DEBUG_DHCPC
	serial0_tx_PGM(PSTR("dhcpc: rx DHCP packet!\r\n"));
#endif

	/* retrieve options, including packet type */
	/* each option consists of a type, followed by possibly a length. */
	while (!nomoreoptions) {
		/* read us the type, and switch on that */
		if (encx24j600_rx(&option_code,sizeof(uint8_t)) != sizeof(uint8_t)) {
			serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option code]\r\n"));
			return E_NOTHING;
		}
		switch (option_code) {
			case 0:
				/* pad option */
				/* this has no payload, just do another loop */
				break;
			case 255:
				/* no more options! */
				nomoreoptions = 1;
				break;
			case 1:
				/* Subnet mask */
				if (encx24j600_rx(&option_len,sizeof(uint8_t)) != sizeof(uint8_t)) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option len]\r\n"));
					return E_NOTHING;
				}
				if (option_len != 4) {
					serial0_tx_PGM(PSTR("dhcpc: Option 1 has incorrect length, discarding\r\n"));
					return E_NOTHING;
				}
				if (encx24j600_rx(dhcp_mask,4) != 4) {
					/* bad! */
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [dhcp_type]\r\n"));
					return E_NOTHING;
				}
				break;
			case 3:
				/* Router (ie, gateway off this LAN), we take the first one always */
				if (encx24j600_rx(&option_len,sizeof(uint8_t)) != sizeof(uint8_t)) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding\r\n"));
					return E_NOTHING;
				}
				if (option_len < 4 || option_len % 4 != 0) {
					serial0_tx_PGM(PSTR("dhcpc: Option 3 has incorrect length, discarding\r\n"));
					return E_NOTHING;
				}
				if (encx24j600_rx(dhcp_gw,4) != 4) {
					/* bad! */
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding\r\n"));
					return E_NOTHING;
				}
				if (option_len > 4) {
					uint8_t discard[4], n;
					/* need to discard the remaining router chunks */
					for (n = 0; n < (option_len-4);n+=4) {
						if (encx24j600_rx(discard,4) != 4) {
							serial0_tx_PGM(PSTR("dhcpc: short frame, discarding\r\n"));
						}
					}
				}
				break;
			case 53:
				/* DHCP Message Type */
				if (encx24j600_rx(&option_len,sizeof(uint8_t)) != sizeof(uint8_t)) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option len]\r\n"));
					return E_NOTHING;
				}
				if (option_len != 1) {
					serial0_tx_PGM(PSTR("dhcpc: Option 53 has incorrect length, discarding\r\n"));
					return E_NOTHING;
				}
				if (encx24j600_rx(&dhcp_type,sizeof(uint8_t)) != sizeof(uint8_t)) {
					/* bad! */
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [dhcp_type]\r\n"));
					return E_NOTHING;
				}
				break;
			case 54:
				/* Server Identifier */
				if (encx24j600_rx(&option_len,sizeof(uint8_t)) != sizeof(uint8_t)) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option len]\r\n"));
					return E_NOTHING;
				}
				if (option_len != 4) {
					serial0_tx_PGM(PSTR("dhcpc: Option 54 has incorrect length, discarding\r\n"));
					return E_NOTHING;
				}
				if (encx24j600_rx(dhcp_si,4) != 4) {
					/* bad! */
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [dhcp_type]\r\n"));
					return E_NOTHING;
				}
				break;
			case 58:
				/* Renewal (T1) value, 32-bit in seconds */
				if (encx24j600_rx(&option_len,sizeof(uint8_t)) != sizeof(uint8_t)) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option len]\r\n"));
					return E_NOTHING;
				}
				if (option_len != 4) {
					serial0_tx_PGM(PSTR("dhcpc: Option 58 has incorrect length, discarding\r\n"));
					return E_NOTHING;
				}
				if (encx24j600_rx(&dhcp_renewal,4) != 4) {
					/* bad! */
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [dhcp_type]\r\n"));
					return E_NOTHING;
				}
				break;
			case 59:
				/* Rebind (T2) value, 32-bit in seconds */
				if (encx24j600_rx(&option_len,sizeof(uint8_t)) != sizeof(uint8_t)) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option len]\r\n"));
					return E_NOTHING;
				}
				if (option_len != 4) {
					serial0_tx_PGM(PSTR("dhcpc: Option 59 has incorrect length, discarding\r\n"));
					return E_NOTHING;
				}
				if (encx24j600_rx(&dhcp_rebind,4) != 4) {
					/* bad! */
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [dhcp_type]\r\n"));
					return E_NOTHING;
				}
				break;
			default:
				/* unknown option, they all have len following them so gobble up the option */
				if (encx24j600_rx(&option_len,sizeof(uint8_t)) != sizeof(uint8_t)) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option len unknown]\r\n"));
					return E_NOTHING;
				}
				/* then skip that many packets */
				if (encx24j600_rx_skip(option_len) != option_len) {
					serial0_tx_PGM(PSTR("dhcpc: short frame, discarding [option skip unknown]\r\n"));
					return E_NOTHING;
				}
				break;
		}
	}

	/* we now have the type of packet, so handle it according to state */
	switch (state) {
		case DHCPC_SELECTING:
			/* expecting a DHCPOFFER message, matching our last xid */
			if (dhcp_type != DHCPOFFER) {
				return E_NOTHING;
			}
			/* check if it's the same xid as we last sent, if not, we don't care */
			if (last_xid != (((uint32_t)pkt.xid[0] << 24) | ((uint32_t)pkt.xid[1] << 16) | ((uint32_t)pkt.xid[2] << 8) | ((uint32_t)pkt.xid[3]))) {
				serial0_tx_PGM(PSTR("dhcpc: bogus xid, discarding\r\n"));
				return E_NOTHING;
			}
			/* if we don't have a server identifier option from the inbound packet, discard it */
			if (dhcp_si[0] == 0x0 && dhcp_si[1] == 0 && dhcp_si[2] == 0 && dhcp_si[3] == 0) {
				serial0_tx_PGM(PSTR("dhcpc: bogus packet, no server identifier option\r\n"));
				return E_NOTHING;
			}
			/* copy the address into our sending address from now on */
			//memcpy(ip_addr,pkt.yiaddr,4);
			/* okay, let's accept this offer since we'll just take whatever one we get first */
			{
				dhcp_rx_header_t tx;
				uint8_t buf[6]; /* used to add options */

#ifdef DEBUG_DHCPC
				serial0_tx_PGM(PSTR("dhcpc: accepting offer\r\n"));
#endif

				/* FIXME: this replies to the DHCPOFFER we got, but not by broadcast. Other servers don't know we've picked it up */

				memset(&tx, 0, sizeof(dhcp_rx_header_t)); /* start with empty packet */

				/* populate the base fields */
				tx.op = BOOTREQUEST;
				tx.htype = 1; /* Ethernet */
				tx.hlen = 6; /* Ethernet Hardware addresses are 6 octets long */

				tx.xid[0] = last_xid >> 24 & 0xff;
				tx.xid[1] = last_xid >> 16 & 0xff;
				tx.xid[2] = last_xid >> 8 & 0xff;
				tx.xid[3] = last_xid & 0xff;
				tx.secs[0] = secs & 0xff; /* DHCPREQUEST needs the same secs as we sent on DHCPDISCOVER */
				tx.secs[1] = secs >> 8;
				tx.flags[0] = 0x80; /* broadcast bit set, not sure we need this but anyway */
				memcpy(tx.chaddr, eth_mac, 6); /* our hardware address */
				//memcpy(tx.ciaddr, pkt.yiaddr, 4); /* the address we were provided */
				tx.magic[0] = 0x63; /* magic number of magic */
				tx.magic[1] = 0x82;
				tx.magic[2] = 0x53;
				tx.magic[3] = 0x63;

				/* build options, in reverse order */
				/* end option */
				buf[0] = 0xff;
				if (!encx24j600_tx_prepend(slot,(uint8_t*)buf,1)) {
					serial0_tx_PGM(PSTR("dhcpc: failed to send DHCPREQUEST\r\n"));
					return E_NOTHING;
				}
				/* option 50, requested IP address */
				buf[0] = 50;
				buf[1] = 4; /* an IP address */
				memcpy(buf+2,pkt.yiaddr,4);
				if (!encx24j600_tx_prepend(slot,(uint8_t*)buf,6)) {
					serial0_tx_PGM(PSTR("dhcpc: failed to send DHCPREQUEST\r\n"));
					return E_NOTHING;
				}
				/* option 54, server identifer */
				buf[0] = 54;
				buf[1] = 4;
				memcpy(buf+2,dhcp_si,4);
				if (!encx24j600_tx_prepend(slot,(uint8_t*)buf,6)) {
					serial0_tx_PGM(PSTR("dhcpc: failed to send DHCPREQUEST\r\n"));
					return E_NOTHING;
				}
				/* option 53, dhcp request type */
				buf[0] = 53;
				buf[1] = 1;
				buf[2] = DHCPREQUEST;
				if (!encx24j600_tx_prepend(slot,(uint8_t*)buf,3)) {
					serial0_tx_PGM(PSTR("dhcpc: failed to send DHCPREQUEST\r\n"));
					return E_NOTHING;
				}

				/* put the rest of the header in the output slot */
				if (!encx24j600_tx_prepend(slot,(uint8_t *)&tx,sizeof(dhcp_rx_header_t))) {
					serial0_tx_PGM(PSTR("dhcpc: failed to send DHCPREQUEST\r\n"));
					return E_NOTHING;
				}
				/* change state and then return our packet size */
				state = DHCPC_REQUESTING;
				timer = 30; /* wait 30 seconds before giving up */
				return E_REPLY_BROAD; /* please send this packet to broadcast */
			}
			break;
		case DHCPC_REQUESTING:
			/* at this point we have a good address and we should be expecting a DHCPACK mesage back. Offers should just be rejected */
			if (dhcp_type == DHCPACK) {
				/* we're confirmed to be using this address, copy it into the local stack and mark us configured. */
				memcpy(ip_addr,pkt.yiaddr,4);
				set_flag(flag_ip,FLAG_IP_CONFIG); /* we are configured! */
				/* set the point we want to consider renewing, and rebinding */
				if (dhcp_renewal) {
					renewal = clock_s + swap16(dhcp_renewal);
				}
				if (dhcp_rebind) {
					rebind = clock_s + swap16(dhcp_rebind);
				}
				/* copy in the discovered mask/gateways */
				memcpy(ip_mask,dhcp_mask,4);
				memcpy(ip_gateway,dhcp_gw,4);
				/* that's pretty much it, rest is handled by dhcpc_tick */
				state = DHCPC_BOUND;
				serial0_tx_PGM(PSTR("dhcpc: ip address is "));
				ip_printaddr(ip_addr);
				serial0_tx_cout('/');
				ip_printaddr(ip_mask);
				serial0_tx_PGM(PSTR(" gw "));
				ip_printaddr(ip_gateway);
				serial0_tx_PGM(PSTR(" [renew in "));
				serial0_tx_dec(swap32(dhcp_renewal));
				serial0_tx_PGM(PSTR("s]\r\n"));
				return E_NOTHING;
			}
			if (dhcp_type == DHCPNAK) {
				/* server somehow allocated our IP address in the mean time! reset back to DHCPC_INIT */
				memset(ip_addr,0,4); /* clear any saved address */
				discoverwait = 0; /* just do it immediate! */
				state = DHCPC_INIT;
				timer = 0;
				return E_NOTHING; /* dhcpc_tick will take care of it from here */
			}
			/* all other kinds of DHCP packets we just ignore */
			break;
	}

	/* FIXME: add handling for DHCPC_RENEWING and DHCPC_REBINDING */

	/* we don't handle DHCPC_BOUND state at all since any incoming packets are just ignored */

	return E_NOTHING;
}

/* handle timed events */

void dhcpc_tick(void) {

	/* don't run the timer if we don't have one */
	if (!(flag_ip & FLAG_IP_DHCP_TIMER)) {
		return;
	}

	/* decrement timer if we have time left to run before doing something */
	if (timer) {
		/* still time to go, come back later */
		timer--;
		return;
	}

	/* what we do depends on the current state */
	switch (state) {
		case DHCPC_INIT:
#ifdef DEBUG_DHCPC
			serial0_tx_PGM(PSTR("dhcpc: DHCPDISCOVER ("));
			serial0_tx_dec(discoverwait+1);
			serial0_tx_PGM(PSTR("s timeout)\r\n"));
#endif
			/* send a UDP packet with a DHCPDISCOVER message */
			{
				dhcp_tx_header_t pkt;
				uint8_t dst[4] = {0xff, 0xff, 0xff, 0xff};

				/* FIXME: most of this is static content we should store in flash and then fill */

				memset(&pkt,0,sizeof(dhcp_tx_header_t)); /* start with empty packet */

				/* populate the base fields */
				pkt.op = BOOTREQUEST;
				pkt.htype = 1; /* Ethernet */
				pkt.hlen = 6; /* Ethernet Hardware addresses are 6 octets long */
				last_xid = random(); /* a .. "random" number */
				pkt.xid[0] = last_xid >> 24 & 0xff;
				pkt.xid[1] = last_xid >> 16 & 0xff;
				pkt.xid[2] = last_xid >> 8 & 0xff;
				pkt.xid[3] = last_xid & 0xff;
				if (!start) {
					start = clock_s;
				}
				secs = (clock_s - start);
				pkt.secs[0] = secs & 0xff; /* how long we've been waiting */
				pkt.secs[1] = secs >> 8;
				pkt.flags[0] = 0x80; /* broadcast bit set, not sure we need this but anyway */ /* unicast maybe? */
				memcpy(pkt.chaddr,eth_mac,6); /* our hardware address */
				pkt.magic[0] = 0x63; /* magic number of magic */
				pkt.magic[1] = 0x82;
				pkt.magic[2] = 0x53;
				pkt.magic[3] = 0x63;
				pkt.option_code = 53;
				pkt.option_len = 1;
				pkt.option_type = DHCPDISCOVER;
				pkt.end = 0xff;

				/* send the UDP packet */
				udp_send(dst,67,68,(uint8_t *)&pkt,sizeof(dhcp_tx_header_t));
			}
			timer = discoverwait; /* send again in n seconds if no response */
			state = DHCPC_SELECTING; /* wait for responses */
			break;
		case DHCPC_SELECTING:
#ifdef DEBUG_DHCPC
			serial0_tx_PGM(PSTR("dhcpc: no offers\r\n"));
#endif
			if (discoverwait < 64) {
				discoverwait = discoverwait << 1;
			}
			state = DHCPC_INIT;
			timer = 0;
			break;
		case DHCPC_REQUESTING:
			/* we should retransmit our request, but instead we'll just assume bad things happened to go back to init */
#ifdef DEBUG_DHCPC
			serial0_tx_PGM(PSTR("dhcpc: no ack to request\r\n"));
#endif
			memset(ip_addr,0,4); /* clear any saved address */
			discoverwait = 2;
			state = DHCPC_INIT;
			timer = 0;
			break;
	}

	return;
}	

/* called when ethernet interface goes link down */
void dhcpc_linkdown(void) {
	return;
} 
